import numpy as np
import random


class ICP:
    @staticmethod
    def svd(points1: np.ndarray, points2: np.ndarray) -> np.ndarray:
        rows = points1.shape[0]
        p1_mean = np.mean(points1, 0)
        p2_mean = np.mean(points2, 0)
        q1 = points1 - p1_mean
        q2 = points2 - p2_mean

        W = np.zeros((3, 3))
        for i in range(rows):
            W += np.reshape(q2[i, :], (3, 1)) @ np.reshape(q1[i, :], (1, 3))

        u, s, vh = np.linalg.svd(W)
        R = u @ vh
        t = p2_mean - R @ p1_mean

        return np.vstack((np.vstack((R.T, t)).T, [0, 0, 0, 1]))

    @staticmethod
    def ransac_svd(points1: np.ndarray, points2: np.ndarray, error=0.2, max_iters=100) -> np.ndarray:
        size = points1.shape[0]

        result_best = np.eye(4)

        if size == 0:
            return result_best

        pretotal = 0
        ratio = 0.75
        P = 0.99
        iter = 0
        while iter < max_iters:
            sample_index = random.sample(range(size), np.min([size, 10]))
            points_1 = points1[sample_index, :]
            points_2 = points2[sample_index, :]

            Te = ICP.svd(points_1, points_2)
            R = Te[:3, :3]
            t = Te[:3, 3]

            total_inlier = 0

            for i in range(size):
                if np.linalg.norm((R @ points1[i, :] + t) - points2[i, :]) < error:
                    total_inlier += 1

            if total_inlier > pretotal:
                if total_inlier < size:
                    max_iters = np.log(1 - P) / np.log(1 - pow(total_inlier / size, 2))
                pretotal = total_inlier
                result_best = Te.copy()

            if total_inlier >= size * ratio:
                break

            iter += 1

        return result_best
